#!/usr/bin/env python3
"""
交互式命令行聊天应用
使用本地 Ollama 运行的 Qwen 模型
"""

from util.llm import qwen_chat
import sys
import time

def main():
    print("交互式聊天机器人")
    print("输入 'quit' 或 'exit' 退出程序")
    print("-" * 40)
    
    # 设置默认的系统提示
    system_prompt = "你是一个乐于回答各种问题的小助手，你的任务是提供专业、准确、有洞察力的建议。"
    
    # 初始化对话历史
    conversation_history = []
    
    while True:
        try:
            # 获取用户输入
            user_input = input("\n你: ").strip()
            
            # 检查退出命令
            if user_input.lower() in ['quit', 'exit']:
                print("AI: 再见！")
                break
            
            # 跳过空输入
            if not user_input:
                continue
            
            # 构建消息历史
            messages = [{"role": "system", "content": system_prompt}]
            messages.extend(conversation_history)
            messages.append({"role": "user", "content": user_input})
            
            # 调用模型获取流式响应
            response = qwen_chat(
                system_prompt=system_prompt,
                user_prompt=user_input,
                stream=True
            )
            
            # 流式显示响应
            print("AI: ", end="", flush=True)
            full_response = ""
            for chunk in response:
                chunk_content = chunk.choices[0].delta.content or ""
                full_response += chunk_content
                print(chunk_content, end="", flush=True)
                time.sleep(0.02)  # 增加一点延迟以获得更好的视觉效果
            
            print()  # 换行
            
            # 将当前对话轮次添加到历史记录中
            conversation_history.append({"role": "user", "content": user_input})
            conversation_history.append({"role": "assistant", "content": full_response})
            
            # 限制历史记录长度，防止上下文过长
            if len(conversation_history) > 10:  # 保留最近5轮对话
                conversation_history = conversation_history[-10:]
            
        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        except Exception as e:
            print(f"发生错误: {e}")

if __name__ == "__main__":
    main()